package files

import (
	"archive/zip"
	"io"
	"os"
	"path/filepath"
	"strings"
)

// ZipDir 将目录压缩为ZIP文件。
//
// 参数:
//   - srcDir: 要压缩的目录路径
//   - zipFilePath: 生成的ZIP文件路径
//
// 返回值:
//   - error: 压缩过程中出现的错误
//
// 注意:
//   - 会递归压缩目录下所有文件和子目录
//   - 自动使用Deflate压缩算法
func ZipDir(srcDir string, zipFilePath string) error {
	zipfile, err := os.Create(zipFilePath)
	if err != nil {
		return err
	}
	defer zipfile.Close()

	archive := zip.NewWriter(zipfile)
	defer archive.Close()

	filepath.Walk(srcDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return err
		}

		header, err := zip.FileInfoHeader(info)
		if err != nil {
			return err
		}

		header.Name = strings.TrimPrefix(path, filepath.Dir(srcDir)+"/")
		// header.Name = path
		if info.IsDir() {
			header.Name += "/"
		} else {
			header.Method = zip.Deflate
		}

		writer, err := archive.CreateHeader(header)
		if err != nil {
			return err
		}

		if !info.IsDir() {
			file, err := os.Open(path)
			if err != nil {
				return err
			}
			defer file.Close()
			_, err = io.Copy(writer, file)
		}
		return err
	})

	return err
}

// ZipFiles 将多个文件压缩为ZIP文件，并可重命名文件路径。
//
// 参数:
//   - files: 要压缩的文件路径列表
//   - zipFilePath: 生成的ZIP文件路径
//   - oldForm: 要替换的原始路径部分
//   - newForm: 替换后的新路径部分
//
// 返回值:
//   - error: 压缩过程中出现的错误
//
// 注意:
//   - 自动使用Deflate压缩算法
//   - 可以批量替换文件路径前缀
func ZipFiles(files []string, zipFilePath string, oldForm, newForm string) error {
	newZipFile, err := os.Create(zipFilePath)
	if err != nil {
		return err
	}
	defer func() {
		_ = newZipFile.Close()
	}()

	zipWriter := zip.NewWriter(newZipFile)
	defer func() {
		_ = zipWriter.Close()
	}()

	for _, file := range files {
		err = func(file string) error {
			zipFile, err := os.Open(file)
			if err != nil {
				return err
			}
			defer zipFile.Close()

			info, err := zipFile.Stat()
			if err != nil {
				return err
			}

			header, err := zip.FileInfoHeader(info)
			if err != nil {
				return err
			}

			// Replace file path.
			header.Name = strings.Replace(file, oldForm, newForm, -1)

			// Optimize compression
			// Refer: http://golang.org/pkg/archive/zip/#pkg-constants
			header.Method = zip.Deflate

			writer, err := zipWriter.CreateHeader(header)
			if err != nil {
				return err
			}
			if _, err = io.Copy(writer, zipFile); err != nil {
				return err
			}
			return nil
		}(file)
		if err != nil {
			return err
		}
	}
	return nil
}
